package be.android.forap.crypto;

import java.io.InputStream;
import java.io.OutputStream;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

import org.spongycastle.crypto.PBEParametersGenerator;
import org.spongycastle.crypto.digests.SHA512Digest;
import org.spongycastle.crypto.engines.AESEngine;
import org.spongycastle.crypto.generators.PKCS12ParametersGenerator;
import org.spongycastle.crypto.modes.CBCBlockCipher;
import org.spongycastle.crypto.paddings.PKCS7Padding;
import org.spongycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.spongycastle.crypto.params.ParametersWithIV;

public class AdvancedEncryptionStandard {
	
	private static final String SALT_DIGEST = "SHA1PRNG";
	private static final int SALT_LENGTH = 512/8;
	private final int allowedKeyLength = 128;
	private final int allowedKeyLength2 = 256;
	
	private int keySize;

	public AdvancedEncryptionStandard(int kSize){
	//public AdvancedEncryptionStandard(int kSize) throws InvalidKeyException{
		if(keySizeIsAllowed(kSize))
			this.keySize = kSize;
		else
			this.keySize = 512;
			//throw new InvalidKeyException("Invalid Key Size");
		
	}
	
	private boolean keySizeIsAllowed(int kSize) {
		return ((kSize==allowedKeyLength)||(kSize==allowedKeyLength2));
	}

	
	/**
	 * This standard also does not prescribe a length for the salt. 
	 * Ideally, the salt is as long as the output of the hash function being used, 
	 * and consists of completely random bits.
	 * ftp://ftp.rsasecurity.com/pub/pkcs/pkcs-12/pkcs-12v1.pdf
	 */
	String salt;
	
	private byte[] generateSalt() throws NoSuchAlgorithmException {
	  SecureRandom random = SecureRandom.getInstance(SALT_DIGEST);
	  byte[] salt = new byte[SALT_LENGTH];
	  random.nextBytes(salt);
	  return salt;
	}
	
	/**
	 * The iteration count is recommended to be 1024 or more
	 * ftp://ftp.rsasecurity.com/pub/pkcs/pkcs-12/pkcs-12v1.pdf
	 */
	private int iterationCount = 1024;
	
	byte[] buf = new byte[1024];
	
	
	public void encrypt(InputStream fin, OutputStream fout, String password) {
	    try {
	    	//LENGTH OF DIGEST
	        //PKCS12ParametersGenerator pGen = new PKCS12ParametersGenerator(new SHA256Digest());
	    	PKCS12ParametersGenerator pGen = new PKCS12ParametersGenerator(new SHA512Digest());
	        char[] passwordChars = password.toCharArray();
	        final byte[] pkcs12PasswordBytes = PBEParametersGenerator.PKCS12PasswordToBytes(passwordChars);
	        pGen.init(pkcs12PasswordBytes, salt.getBytes(), iterationCount);
	        
	        CBCBlockCipher aesCBC = new CBCBlockCipher(new AESEngine());
	        //For block cipher modes of operation, the IV is usually as large as the block size of the cipher:128
	        int ivSize = 128;
	        ParametersWithIV aesCBCParams = (ParametersWithIV) pGen.generateDerivedParameters(256, ivSize);
	        aesCBC.init(true, aesCBCParams);
	        PaddedBufferedBlockCipher aesCipher = new PaddedBufferedBlockCipher(aesCBC, new PKCS7Padding());
	        aesCipher.init(true, aesCBCParams);

	        // Read in the cleartext and write the encrypted bytes to out
	        int numRead = 0;
			while ((numRead = fin.read(buf)) >= 0) {
	            if (numRead == 1024) {
	                byte[] plainTemp = new byte[aesCipher.getUpdateOutputSize(numRead)];
	                int offset = aesCipher.processBytes(buf, 0, numRead, plainTemp, 0);
	                final byte[] plain = new byte[offset];
	                System.arraycopy(plainTemp, 0, plain, 0, plain.length);
	                fout.write(plain, 0, plain.length);
	            } else {
	                byte[] plainTemp = new byte[aesCipher.getOutputSize(numRead)];
	                int offset = aesCipher.processBytes(buf, 0, numRead, plainTemp, 0);
	                int last = aesCipher.doFinal(plainTemp, offset);
	                final byte[] plain = new byte[offset + last];
	                System.arraycopy(plainTemp, 0, plain, 0, plain.length);
	                fout.write(plain, 0, plain.length);
	            }
	        }
	        fout.close();
	        fin.close();
	    } catch (Exception e) {
	        e.printStackTrace();
	    }

	}

	public void decrypt(InputStream fin, OutputStream fout, String password) {
	    try {
	        //PKCS12ParametersGenerator pGen = new PKCS12ParametersGenerator(new SHA256Digest());
	    	PKCS12ParametersGenerator pGen = new PKCS12ParametersGenerator(new SHA512Digest());
	        char[] passwordChars = password.toCharArray();
	        final byte[] pkcs12PasswordBytes = PBEParametersGenerator.PKCS12PasswordToBytes(passwordChars);
	        pGen.init(pkcs12PasswordBytes, salt.getBytes(), iterationCount);
	        CBCBlockCipher aesCBC = new CBCBlockCipher(new AESEngine());
	        ParametersWithIV aesCBCParams = (ParametersWithIV) pGen.generateDerivedParameters(256, 128);
	        aesCBC.init(false, aesCBCParams);
	        PaddedBufferedBlockCipher aesCipher = new PaddedBufferedBlockCipher(aesCBC, new PKCS7Padding());
	        aesCipher.init(false, aesCBCParams);

	        // Read in the decrypted bytes and write the cleartext to out
	        int numRead = 0;
	        while ((numRead = fin.read(buf)) >= 0) {
	            if (numRead == 1024) {
	                byte[] plainTemp = new byte[aesCipher.getUpdateOutputSize(numRead)];
	                int offset = aesCipher.processBytes(buf, 0, numRead, plainTemp, 0);
	                // int last = aesCipher.doFinal(plainTemp, offset);
	                final byte[] plain = new byte[offset];
	                System.arraycopy(plainTemp, 0, plain, 0, plain.length);
	                fout.write(plain, 0, plain.length);
	            } else {
	                byte[] plainTemp = new byte[aesCipher.getOutputSize(numRead)];
	                int offset = aesCipher.processBytes(buf, 0, numRead, plainTemp, 0);
	                int last = aesCipher.doFinal(plainTemp, offset);
	                final byte[] plain = new byte[offset + last];
	                System.arraycopy(plainTemp, 0, plain, 0, plain.length);
	                fout.write(plain, 0, plain.length);
	            }
	        }
	        fout.close();
	        fin.close();
	    } catch (Exception e) {
	        e.printStackTrace();
	    }
	}


}
